import neptune.new as neptune

run = neptune.init_run(
    project="shao1chuan/log入门例子",
    api_token="eyJhcGlfYWRkcmVzcyI6Imh0dHBzOi8vYXBwLm5lcHR1bmUuYWkiLCJhcGlfdXJsIjoiaHR0cHM6Ly9hcHAubmVwdHVuZS5haSIsImFwaV9rZXkiOiIxNGEwOWZhYS1iM2IwLTQ3ZjItYmE2MS02ZjVjOTNjMTcyZDcifQ==",
)  # your credentials

params = {"learning_rate": 0.001, "optimizer": "Adam"}
run["parameters"] = params

import os

import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
from pytorch_lightning import Trainer
from neptune.new.types import File


PATH_DATASETS = os.environ.get("PATH_DATASETS", "../../data/")
image_path = os.environ.get("PATH_DATASETS", "../../data/dog&cat/")
AVAIL_GPUS = min(1, torch.cuda.device_count())
BATCH_SIZE = 256 if AVAIL_GPUS else 64

class Model(pl.LightningModule):
    def __init__(self, layer_size=784):
        super().__init__()
        self.save_hyperparameters()
        self.l1 = torch.nn.Linear(layer_size, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_nb):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        run["train/loss"].append(loss)
        return loss

    def validation_step(self, batch, batch_nb):
        x, y = batch
        y_hat = self.forward(x)
        loss = F.cross_entropy(y_hat, y)
        run["val/loss"].append(loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


# Init our model
model = Model()

run = neptune.init_run(
    project="shao1chuan/log入门例子",
    api_token="eyJhcGlfYWRkcmVzcyI6Imh0dHBzOi8vYXBwLm5lcHR1bmUuYWkiLCJhcGlfdXJsIjoiaHR0cHM6Ly9hcHAubmVwdHVuZS5haSIsImFwaV9rZXkiOiIxNGEwOWZhYS1iM2IwLTQ3ZjItYmE2MS02ZjVjOTNjMTcyZDcifQ==",
)  # your credentials


# Init DataLoader from MNIST Dataset
train_ds = MNIST(
    PATH_DATASETS, train=True, download=True, transform=transforms.ToTensor()
)
train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE)

run["val/conf_matrix"].upload("../data/dog&cat/1.jpg")

eval_ds = MNIST(
    PATH_DATASETS, train=False, download=True, transform=transforms.ToTensor()
)
eval_loader = DataLoader(train_ds, batch_size=BATCH_SIZE)

# Initialize a trainer
trainer = Trainer(gpus=AVAIL_GPUS, max_epochs=3)

# Train the model ⚡
trainer.fit(model, train_loader, eval_loader)

# run["eval/f1_score"] = 0.66

run.stop()